package com.only.common.util.ssh;

import cn.hutool.core.date.DateUtil;
import cn.hutool.core.util.StrUtil;
import com.only.common.util.PathUtil;
import lombok.Data;
import org.apache.sshd.client.SshClient;
import org.apache.sshd.client.channel.ChannelExec;
import org.apache.sshd.client.session.ClientSession;
import org.apache.sshd.sftp.client.SftpClientFactory;
import org.apache.sshd.sftp.client.fs.SftpFileSystem;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.io.*;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.*;
import java.util.concurrent.TimeUnit;

import static java.nio.file.StandardCopyOption.*;

/**
 * @author wentao.zou
 * @date 2023/12/06 16:47
 */
public class SshUtil {
    private static final Logger logger = LoggerFactory.getLogger(SshUtil.class);

    /**
     * 上传文件(如果文件已存在，会先删除掉，使用前做好备份）
     */
    public static UploadFileResult uploadFile(
        SshBaseParams sshBaseParams,
        String filePath,
        String fileName,
        String fileContent,
        boolean needBackups
    ) {
        SshClientSession sshClientSession = new SshClientSession(sshBaseParams);
        try {
            sshClientSession.init();
            return SshUtil.uploadFileBySession(sshClientSession, filePath, fileName, fileContent, needBackups);
        } catch (Exception ex) {
            logger.error("无法上传文件", ex);
            return new UploadFileResult(StrUtil.format("上传文件失败，异常：{}", ex.getMessage()), "", PathUtil.add(filePath, fileName));
        } finally {
            sshClientSession.close();
        }
    }

    /**
     * 将本地文件上传到服务器
     * 如果文件已经存在会被覆盖，并且没有备份
     * @param sshBaseParams
     * @param srcFileName 本地文件（绝对文件名字）
     * @param remoteFileName 远程文件 （绝对文件名字）
     */
    public static UploadFileResult uploadFile(SshBaseParams sshBaseParams, String srcFileName, String remoteFileName) {
        SshClientSession sshClientSession = new SshClientSession(sshBaseParams);
        try {
            sshClientSession.init();
            // 上传文件
            SftpFileSystem fs = SftpClientFactory.instance().createSftpFileSystem(sshClientSession.getSftpSession());
            fs.setWriteBufferSize(32 * 1024);
            fs.setReadBufferSize(32 * 1024);

            Path remoteRoot = fs.getDefaultDir().resolve(PathUtil.getParent(remoteFileName));

            // 检查并创建目录
            if (!Files.exists(remoteRoot)) {
                Files.createDirectories(remoteRoot);
            }

            Path remoteFile = fs.getDefaultDir().resolve(remoteFileName);
            Path srcFile = new File(srcFileName).toPath();

            Files.copy(srcFile, remoteFile, REPLACE_EXISTING);
            fs.close();

            return new UploadFileResult(srcFileName, remoteFileName);
        } catch (Exception ex) {
            logger.error("无法上传文件：" + remoteFileName, ex);
            return new UploadFileResult( StrUtil.format("上传文件失败，异常：{}", ex.getMessage()), srcFileName, remoteFileName);
        } finally {
            sshClientSession.close();
        }
    }

    /**
     * 下载文件，将远程服务器的文件 remoteFileName 下载到本地 localFileName
     * @param sshBaseParams
     * @param localFileName 本地文件名字
     * @param remoteFileName 远程文件文件
     * @return
     */
    public static DownloadFileResult downloadFile(SshBaseParams sshBaseParams, String localFileName, String remoteFileName) {
        SshClientSession sshClientSession = new SshClientSession(sshBaseParams);
        try {
            sshClientSession.init();
            // 下载文件
            SftpFileSystem fs = SftpClientFactory.instance().createSftpFileSystem(sshClientSession.getSftpSession());
            Path localRoot = new File(PathUtil.getParent(localFileName)).toPath();

            // 检查并创建目录
            if (!Files.exists(localRoot)) {
                Files.createDirectories(localRoot);
            }

            Path remoteFile = fs.getDefaultDir().resolve(remoteFileName);
            Path srcFile = new File(localFileName).toPath();

            Files.copy(remoteFile, srcFile, REPLACE_EXISTING);
            fs.close();

            return new DownloadFileResult(localFileName, remoteFileName);
        } catch (Exception ex) {
            logger.error("无法下载文件：" + remoteFileName, ex);
            return new DownloadFileResult(StrUtil.format("无法下载文件！异常：{}", ex.getMessage()), localFileName, remoteFileName);
        } finally {
            sshClientSession.close();
        }
    }

    /**
     * 通用执行命令方法
     */
    public static CmdExecuteResult executeCmd(SshBaseParams sshBaseParams, String command) {
        SshClientSession sshClientSession = new SshClientSession(sshBaseParams);
        try {
            sshClientSession.init();
            return executeCmdBySession(sshClientSession, command);
        } catch (Exception ex) {
            logger.error("执行命令异常！, 指令:" + command, ex);
            return new CmdExecuteResult(false, -1, ex.getMessage());
        } finally {
            sshClientSession.close();
        }
    }

    /**
     * 多条命令保存成文件再一次执行（执行完成删除命令文件）
     * @param fileId 文件的唯一ID，将以这个ID为文件名字上传指令脚本到服务器
     */
    public static CmdExecuteResult executeCmdByFile(SshBaseParams sshBaseParams, String fileId, SshCommand cmd) {
        String remotePath = "/tmp/dcp-shell";
        String fileName = StrUtil.format("{}.sh", fileId);
        String absFileName = PathUtil.add(remotePath, fileName);

        SshClientSession sshClientSession = new SshClientSession(sshBaseParams);
        CmdExecuteResult cmdExecuteResult;
        try {
            sshClientSession.init();
            UploadFileResult uploadFile = SshUtil.uploadFileBySession(
                sshClientSession,
                remotePath,
                fileName,
                cmd.buildCommand(),
                false
            );

            if (!uploadFile.isSuccess()) {
                return new CmdExecuteResult(false, -1, "无法上传指令文件");
            }

            String command = StrUtil.format("cd {} && sh {} ", remotePath, absFileName);
            CmdExecuteResult executeResult = SshUtil.executeCmdBySession(sshClientSession, command);
            cmdExecuteResult = new CmdExecuteResult(executeResult.isSuccessful(), executeResult.getStatus(), executeResult.getExecuteLog());
        } catch (Exception ex) {
            logger.error("执行指令文件异常！指令：" + cmd.buildCommand(), ex);
            cmdExecuteResult = new CmdExecuteResult(false, -1, ex.getMessage());
        } finally {
            SshUtil.executeCmdBySession(sshClientSession, StrUtil.format("rm -rf {}", absFileName));// 抛异常时也要删除文件
            sshClientSession.close();
        }
        return cmdExecuteResult;
    }

    /**
     * 以 ‘与’ 的方式执行指令列表
     */
    public static CmdExecuteResult executeCommandsWithAnd(SshBaseParams sshBaseParams, SshCommand command) {
        String cmd = command.buildCommandsByAnd();
        CmdExecuteResult executeResult = SshUtil.executeCmd(sshBaseParams, cmd);
        return new CmdExecuteResult(executeResult.isSuccessful(), executeResult.getStatus(), executeResult.getExecuteLog());
    }

    public static CmdExecuteResultMap executeCommandsToMap(SshBaseParams sshBaseParams, SshCommand command){
        List<String> commands = command.getCommands();
        SshClientSession sshClientSession = new SshClientSession(sshBaseParams);
        try {
            sshClientSession.init();
            HashMap<String, String> hashMap = new HashMap<>();

            for (String cmd : commands){
                long startTime = System.currentTimeMillis();

                CmdExecuteResult result = executeCmdBySession(sshClientSession, cmd);
                hashMap.put(cmd, result.getExecuteLog());

                logger.info("执行命令 {} 花费: {}", cmd, System.currentTimeMillis() - startTime);
            }

            return new CmdExecuteResultMap(true, hashMap);
        } catch (Exception ex) {
            logger.error("执行命令异常！", ex);
            return new CmdExecuteResultMap(false, null);
        } finally {
            sshClientSession.close();
        }
    }

    /**
     * 端口检测
     */
    public static PortCheckResult checkPort(SshBaseParams sshBaseParams, int checkPort) {
        String cmd = StrUtil.format("/usr/sbin/ss -tulpn | grep -w '{}'", checkPort);

        CmdExecuteResult ret = executeCmd(sshBaseParams, cmd);
        if (ret.isSuccessful()) {
            return new PortCheckResult(true, StrUtil.isNotBlank(ret.getExecuteLog()));
        } else {
            return new PortCheckResult(false,false, ret.getExecuteLog());
        }
    }

    /**
     * 获取服务器上所有未使用的端口
     * @param sshBaseParams
     * @return
     */
    public static UsedPortsResult listUsedPorts(SshBaseParams sshBaseParams) {
        SshCommand cmd = new SshCommand("/usr/sbin/ss -tulpn | awk '{print $5}'");

        SshClientSession sshClientSession = new SshClientSession(sshBaseParams);
        try {
            sshClientSession.init();
            String ret = sshClientSession.executeRemoteCommand(cmd.getRawCommands());
            return UsedPortsResult.parsingPorts(ret);
        } catch (Exception ex) {
            logger.error("获取服务器上所有未使用的端口！", ex);
            return new UsedPortsResult(ex.getMessage());
        } finally {
            sshClientSession.close();
        }
    }

    private static CmdExecuteResult executeCmdBySession(SshClientSession sshClientSession, String command){
        StringBuilder executeLog = new StringBuilder();

        // 执行远程命令
        try (ChannelExec channel = sshClientSession.getCmdSession().createExecChannel(command)) {
            channel.open().verify(60, TimeUnit.SECONDS);

            InputStream in = channel.getInvertedOut();
            InputStreamReader reader = new InputStreamReader(in);
            BufferedReader bufferedReader = new BufferedReader(reader);

            String line;
            while ((line = bufferedReader.readLine()) != null) {
                executeLog.append(line).append("\n");
            }
            
            int status = channel.getExitStatus();
            // 关闭资源
            bufferedReader.close();
            reader.close();
            in.close();

            return new CmdExecuteResult(true, status, executeLog.toString());
        } catch (Exception ex) {
            logger.error("执行指令异常！", ex);
            return new CmdExecuteResult(false, -1, "无法链接服务器，无法执行指令！");
        }
    }

    private static UploadFileResult uploadFileBySession(
        SshClientSession sshClientSession,
        String remoteFilePath,
        String fileName,
        String fileContent,
        boolean needBackups
    ) {
        try {
            // 上传文件
            SftpFileSystem fs = SftpClientFactory.instance().createSftpFileSystem(sshClientSession.getSftpSession());
            Path remoteRoot = fs.getDefaultDir().resolve(remoteFilePath);

            // 检查并创建目录
            if (!Files.exists(remoteRoot)) {
                Files.createDirectories(remoteRoot);
            }

            Path remoteFile = remoteRoot.resolve(PathUtil.add(remoteFilePath, fileName));

            // 如果需要，先备份文件
            if (Files.exists(remoteFile) && needBackups) {
                Path bakPath = remoteRoot.resolve(".dcp_bak");
                if (!Files.exists(bakPath)) {
                    Files.createDirectories(bakPath);
                }

                String time = DateUtil.format(new Date(), "yyyyMMddHHmmss.SSS");
                String backFileName = StrUtil.format("{}-{}", fileName, time);

                Path absBackFileName = bakPath.resolve(backFileName);
                Files.copy(remoteFile, absBackFileName, REPLACE_EXISTING);
            }

            // 写入内容，覆盖文件
            byte[] bytes = fileContent.getBytes();
            InputStream inputStream = new ByteArrayInputStream(bytes);
            Files.copy(inputStream, remoteFile, REPLACE_EXISTING);

            fs.close();
            return new UploadFileResult("", PathUtil.add(remoteFilePath, fileName));
        } catch (IOException e) {
            // 连接失败的异常
            logger.error("文件上传异常！", e);
            return new UploadFileResult("无法上传文件！", "", PathUtil.add(remoteFilePath, fileName));
        }
    }

    @Data
    private static class SshClientSession {
        private static final Logger logger = LoggerFactory.getLogger(SshClientSession.class);

        // client 是全局单例
        private volatile static SshClient client;
        // 运行指令的 session
        private ClientSession cmdSession;
        // 上传文件的 session
        private ClientSession sftpSession;
        private final SshBaseParams sshBaseParams;

        public SshClientSession(SshBaseParams sshBaseParams) {
            this.sshBaseParams = sshBaseParams;
        }

        public void init() throws Exception {
            synchronized (this) {
                if (!initialized()) {
                    initSshClient();
                    initCmdSession();
                    initSftpSession(); // 如果端口相同，用相同的 session
                }
            }
        }

        /**
         * 直接执行远程的指令, 只能执行一些比较简单的指令行
         * @param cmd 指令
         * @return
         */
        public String executeRemoteCommand(String cmd) throws Exception {
            if (initialized()) {
                try {
                    long start = System.currentTimeMillis();
                    String ret = this.cmdSession.executeRemoteCommand(cmd);
                    logger.info(StrUtil.format("链接 {}\n执行指令 {}\n用时：{}ms", sshBaseParams.host, cmd, System.currentTimeMillis() - start));
                    return ret;
                } catch (Exception ex) {
                    logger.error(StrUtil.format("无法执行指令:\n{}\n异常信息：{}", cmd, ex.getMessage()), ex);
                    throw new Exception("执行指令失败！");
                }
            } else {
                throw new Exception("服务器 session 未初始化");
            }
        }

        public void close() {
            closeSession();
        }

        private void initCmdSession() throws Exception {
            try {
                ClientSession session = this.client.connect(
                    sshBaseParams.getUsername(),
                    sshBaseParams.getHost(),
                    sshBaseParams.getPort()
                ).verify(sshBaseParams.getTimeOut()).getSession();

                session.addPasswordIdentity(sshBaseParams.getPassword());
                session.auth().verify(sshBaseParams.getTimeOut());

                this.cmdSession = session;
            } catch (Exception ex) {
                logger.error("无法初始化 CmdSession", ex);
                throw new Exception(StrUtil.format("无法连接服务器或者账号密码错误！服务器：{}", sshBaseParams.host));
            }
        }

        private void initSftpSession() throws Exception {
            if (Objects.equals(sshBaseParams.sftpPort, sshBaseParams.getPort())) {
                sftpSession = cmdSession;
            } else {
                try {
                    long start = System.currentTimeMillis();
                    ClientSession session = this.client.connect(
                        sshBaseParams.getUsername(),
                        sshBaseParams.getHost(),
                        sshBaseParams.getSftpPort()
                    ).verify(sshBaseParams.getTimeOut()).getSession();
                    logger.info(StrUtil.format("链接 {}，创建 sftp session 耗时 {} ms", sshBaseParams.host, System.currentTimeMillis() - start));

                    start = System.currentTimeMillis();
                    session.addPasswordIdentity(sshBaseParams.getPassword());
                    session.auth().verify(sshBaseParams.getTimeOut());
                    logger.info(StrUtil.format("链接 {}，sftp 密码认证耗时 {} ms", sshBaseParams.host, System.currentTimeMillis() - start));

                    this.sftpSession = session;
                } catch (Exception ex) {
                    logger.error("无法连接服务器或者账号密码错误！", ex);
                    throw new Exception(StrUtil.format("无法连接服务器或者账号密码错误！服务器：{}", sshBaseParams.host));
                }
            }
        }

        private void initSshClient() throws Exception {
            if (client == null) {
                synchronized(SshClientSession.class) {
                    if (client == null) {
                        long start = System.currentTimeMillis();
                        SshClient client = SshClient.setUpDefaultClient();
                        logger.info(StrUtil.format("链接 {}，创建 client 耗时 {} ms", sshBaseParams.host, System.currentTimeMillis() - start));

                        try {
                            start = System.currentTimeMillis();
                            client.start();
                            logger.info(StrUtil.format("链接 {}，启动 client 耗时 {} ms", sshBaseParams.host, System.currentTimeMillis() - start));
                            this.client = client;
                        } catch (Exception ex) {
                            throw new Exception("无法连接服务器！");
                        }
                    }
                }
            }
        }

        private boolean initialized() {
            return client != null && cmdSession != null;
        }

        private void closeSession() {
            if (this.cmdSession != null) {
                try {
                    this.cmdSession.close();
                } catch (Exception ignored) { }
            }

            if (this.sftpSession != null) {
                try {
                    this.sftpSession.close();
                } catch (Exception ignored) { }
            }
        }

        private void closeSshClient() {
            if (this.client != null) {
                try {
                    this.client.close();
                } catch (Exception ignored) {}
            }
        }
    }

    @Data
    public static class SshBaseParams {
        /**
         * ip地址
         */
        private String host;

        /**
         * 连接服务器的端口
         */
        private Integer port = 22;

        /**
         * 用于 sftp 的端口
         */
        private Integer sftpPort = 22;

        /**
         * 登录用户名
         */
        private String username;

        /**
         * 密码
         */
        private String password;

        /**
         * 超时时间，定位毫秒
         */
        private long timeOut = 12000L;

        /**
         * 构造 ssh 参数
         * @param host ip
         * @param sshPort 端口，默认 ssh 端口 和 sftp 端口一致
         * @param username 用户
         * @param password 密码
         */
        public SshBaseParams(String host, int sshPort, String username, String password) {
            this.host = host;
            this.port = sshPort;
            this.sftpPort = sshPort;
            this.username = username;
            this.password = password;
        }

        public SshBaseParams setSftpPort(int sftpPort) {
            this.sftpPort = sftpPort;
            return this;
        }
    }
}
